import tkinter as tk
from tkinter import ttk
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib.figure import Figure
import numpy as np
import cv2
import torchvision.transforms as transforms
import torch
from CNN import SimpleCNN  # 确保要有这个文件和对应的模型
import mss
import time
import threading
from queue import Queue

# 常量
STD_LIST = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, "+", '-', ".", '*', "(", ")", '/']
SCALE_FACTOR = 0.5  # 缩放因子

# 初始化模型和转换
def initialize_model():
    model = SimpleCNN()
    model.load_state_dict(torch.load('mnist_cnn.pth'))
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    model.eval()
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    return model, transform, device

model, transform, device = initialize_model()

class LatexRenderer:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("LaTeX 实时渲染")
        self.root.geometry('1000x500')
        self.root.configure(bg='lightgray')

        self.fig = Figure(figsize=(14, 2))
        self.ax = self.fig.add_subplot(111)
        self.ax.axis('off')
        self.canvas = FigureCanvasTkAgg(self.fig, master=self.root)
        self.canvas.get_tk_widget().pack(padx=10, pady=10)

        self.status_label = ttk.Label(self.root, text="准备就绪", background='lightgray', font=('Arial', 12))
        self.status_label.pack(pady=10)

        self.latex_formula = None
        self.update_latex()

        self.entry = tk.Entry(self.root, font=('Arial', 12))
        self.entry.pack(pady=10)
        self.entry_button = ttk.Button(self.root, text="更新公式", command=self.manual_update)
        self.entry_button.pack(pady=10)

    def update_latex(self):
        if self.latex_formula:
            self.ax.clear()
            self.ax.axis('off')
            self.ax.text(0.5, 0.5, f"${self.latex_formula}$", fontsize=20, ha='center', va='center')
            self.canvas.draw()
        self.root.after(100, self.update_latex)

    def render_latex(self, latex_formula):
        self.latex_formula = latex_formula

    def update_status(self, status_text):
        self.status_label.config(text=status_text)

    def manual_update(self):
        formula = self.entry.get()
        try:
            self.render_latex(formula + " = " + str(eval(formula)))
            self.update_status("公式更新成功")
        except Exception as e:
            self.render_latex("Error! " + formula)
            self.update_status(f"错误: {str(e)}")

renderer = LatexRenderer()

def accessBinary(img, threshold=128):
    img = 255 - img  # 反转颜色
    kernel = np.ones((3, 3), np.uint8)
    img = cv2.dilate(img, kernel, iterations=1)
    _, img = cv2.threshold(img, threshold, 0, cv2.THRESH_TOZERO)
    return img

def showResults(img, borders, results=None):
    if borders is None:
        return
    res = ""  # 最后的式子
    sorted_results = sorted(zip(borders, results), key=lambda x: x[0][0][0])
    for i, (border, result) in enumerate(sorted_results):
        if result is not None:
            res += str(STD_LIST[result])
    try:
        res = res[:len(res) - 3]
        renderer.render_latex(res + " = " + str(eval(res)))  # 调用 render_latex 函数
        renderer.update_status("检测到公式")
    except Exception as e:
        renderer.render_latex("Error! " + res)  # 显示识别到的错误结果
        renderer.update_status(f"错误: {str(e)}")

def transMNIST(img, borders, size=(28, 28)):
    imgData = np.zeros((len(borders), size[0], size[0], 1), dtype='uint8')
    img = accessBinary(img)
    for i, border in enumerate(borders):
        borderImg = img[border[0][1]:border[1][1], border[0][0]:border[1][0]]
        extendPixel = (max(borderImg.shape) - min(borderImg.shape)) // 2
        targetImg = cv2.copyMakeBorder(borderImg, 7, 7, extendPixel + 7, extendPixel + 7, cv2.BORDER_CONSTANT)
        targetImg = cv2.resize(targetImg, size)
        targetImg = np.expand_dims(targetImg, axis=-1)
        imgData[i] = targetImg
    return imgData

def findBorderContours(img, maxArea=50):
    img = accessBinary(img)
    contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    borders = []
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        if w * h > maxArea:
            border = [(x, y), (x + w, y + h)]
            borders.append(border)
    return borders

def preprocess_image(image):
    image = transform(image)
    image = image.unsqueeze(0)
    return image

def predict(model, image):
    input_tensor = preprocess_image(image).to(device)
    with torch.no_grad():
        output = model(input_tensor)
        pred_label = output.argmax(dim=1, keepdim=True).item()
        return pred_label

def capture_screen(queue, monitor, stop_event):
    sct = mss.mss()
    while not stop_event.is_set():
        screenshot = np.array(sct.grab(monitor))
        queue.put(screenshot)
        time.sleep(0.1)  # 控制帧率

def process_frame(model, frame_queue, stop_event):
    last_recognition_time = time.time()
    recognition_interval = 0.5

    while not stop_event.is_set():
        if not frame_queue.empty():
            screenshot = frame_queue.get()
            screenshot = cv2.cvtColor(screenshot, cv2.COLOR_BGRA2BGR)
            current_time = time.time()

            if current_time - last_recognition_time >= recognition_interval:
                gray_screenshot = cv2.cvtColor(screenshot, cv2.COLOR_BGR2GRAY)
                borders = findBorderContours(gray_screenshot)
                if borders:
                    imgData = transMNIST(gray_screenshot, borders)
                    recognize_results = []
                    for i in range(len(borders)):
                        result = predict(model, imgData[i])
                        recognize_results.append(result)

                    sorted_indexes = sorted(range(len(borders)), key=lambda i: borders[i][0][0])
                    sorted_borders = [borders[i] for i in sorted_indexes]
                    sorted_results = [recognize_results[i] for i in sorted_indexes]

                    showResults(screenshot, sorted_borders, sorted_results)
                else:
                    renderer.update_status("未检测到公式")

                last_recognition_time = current_time
            else:
                showResults(screenshot, borders=None)

def real_time_detection():
    frame_queue = Queue()
    monitor = {"top": 54, "left": 8, "width": 1895, "height": 962}  # 修改为需要的显示器区域

    stop_event = threading.Event()

    screen_thread = threading.Thread(target=capture_screen, args=(frame_queue, monitor, stop_event))
    process_thread = threading.Thread(target=process_frame, args=(model, frame_queue, stop_event))

    screen_thread.start()
    process_thread.start()

    def on_closing():
        stop_event.set()
        screen_thread.join()
        process_thread.join()
        renderer.root.quit()
        renderer.root.destroy()

    renderer.root.protocol("WM_DELETE_WINDOW", on_closing)
    renderer.root.mainloop()

real_time_detection()
